
import json
import pydot
import os
import logging
import sys,subprocess
import argparse
import re,utils


def create_dot_from_json(data,dot_dir):
    state_machine_names = data.keys()
    for state_machine_name in state_machine_names:
        dot_path = os.path.join(dot_dir,f"{state_machine_name}.dot")
        
        # 创建有向图
        graph = pydot.Dot(graph_type='digraph', rankdir='LR')
        graph.set_node_defaults(shape='ellipse', style='filled', fillcolor='white', color='black')
        
        # 收集所有状态
        all_states = set()
        
        for key, value in data[state_machine_name].items():
            try:
                transitions = value['transitions']
                for transition in transitions:
                    all_states.add(transition["start_state"])
                    all_states.add(transition["end_state"])

            except KeyError:
                continue

        # 添加所有状态节点
        for state in all_states:
            node = pydot.Node(state)
            graph.add_node(node)

        # 添加转换边
        for key, value in data[state_machine_name].items():
            try:
                transitions = value['transitions']
                for transition in transitions:
                    start_state = transition["start_state"]
                    end_state = transition["end_state"]
                    condition = transition["condition"]
                    explanation = transition.get('condition explanation', '')
                    message_type = transition.get('message_type', '')
                    message_direction = transition.get('message_direction', '')
                    
                    # 构建边标签
                    label_parts = []
                    if message_type and message_type != 'UNKNOWN':
                        label_parts.append(f"Msg: {message_type}")
                    if message_direction and message_direction != 'UNKNOWN':
                        label_parts.append(f"Dir: {message_direction}")
                    if condition:
                        label_parts.append(f"Cond: {condition}")
                    
                    edge_label = "\\n".join(label_parts) if label_parts else condition
                    
                    # 根据消息方向设置边的颜色
                    edge_color = 'black'
                    if message_direction == 'SEND':
                        edge_color = 'blue'
                    elif message_direction == 'RECEIVE':
                        edge_color = 'green'
                    elif message_direction == 'COMPLETE':
                        edge_color = 'red'
                    
                    edge = pydot.Edge(start_state, end_state, label=edge_label, 
                                    color=edge_color, fontsize=8)
                    graph.add_edge(edge)

            except KeyError:
                continue

        # 保存DOT文件
        graph.write(dot_path)
        
        # 生成PDF
        pdf_path = dot_path[:-3] + 'pdf'
        command = ['dot','-Tpdf',dot_path,'-o',pdf_path]
        subprocess.run(command)
        
        print(f"Generated: {dot_path} and {pdf_path}")

def filter_json_by_config(json_data, config):
    filtered_data = json_data.copy()
    state_machine_names = list(json_data.keys())
    for state_machine_name in state_machine_names:
        state_enum = config.get("state_functions", {}).get(state_machine_name, {}).get("state_enum")
        for key, value in list(json_data[state_machine_name].items()):
            transitions = value['transitions']
            filtered_transitions = []
            for transition in transitions:
                start_state = transition["start_state"]
                end_state = transition["end_state"]
                if state_enum and (start_state not in state_enum and end_state not in state_enum):
                    continue
                filtered_transitions.append(transition)
            # 如果没有合法的transition，则删除
            if filtered_transitions:
                if state_machine_name not in filtered_data:
                    filtered_data[state_machine_name] = {}
                filtered_data[state_machine_name][key] = {"transitions": filtered_transitions}
            else:
                if key in filtered_data.get(state_machine_name, {}):
                    del filtered_data[state_machine_name][key]
        # 如果该state_machine_name下所有key都被删光了，则删除state_machine_name
        if not filtered_data.get(state_machine_name):
            del filtered_data[state_machine_name]
    return filtered_data

def create_enhanced_fsm_json(data, output_dir):
    """
    创建增强的状态机JSON文件，包含完整的IO信息
    """
    enhanced_data = {}
    
    for state_machine_name, functions in data.items():
        enhanced_data[state_machine_name] = {
            "states": set(),
            "transitions": [],
            "messages": set()
        }
        
        for func_name, func_data in functions.items():
            if 'transitions' in func_data:
                # 收集状态
                for transition in func_data['transitions']:
                    enhanced_data[state_machine_name]["states"].add(transition["start_state"])
                    enhanced_data[state_machine_name]["states"].add(transition["end_state"])
                    
                    # 收集消息信息
                    if 'message_type' in transition and transition['message_type'] != 'UNKNOWN':
                        enhanced_data[state_machine_name]["messages"].add(transition['message_type'])
                
                # 收集转换
                enhanced_data[state_machine_name]["transitions"].extend(func_data['transitions'])
        
        # 转换为列表格式
        enhanced_data[state_machine_name]["states"] = list(enhanced_data[state_machine_name]["states"])
        enhanced_data[state_machine_name]["messages"] = list(enhanced_data[state_machine_name]["messages"])
    
    # 保存增强的JSON文件
    enhanced_json_path = os.path.join(output_dir, f"{list(data.keys())[0]}_enhanced_io.json")
    with open(enhanced_json_path, 'w', encoding='utf-8') as f:
        json.dump(enhanced_data, f, indent=2, ensure_ascii=False)
    
    print(f"Generated enhanced FSM JSON: {enhanced_json_path}")
    return enhanced_data


def main(input,output,config):
    os.makedirs(output,exist_ok=True)
    with open(input, 'r', encoding='utf-8') as f:
        json_data = json.load(f)
    
    # json_data = filter_json_by_config(json_data, config)

    # 创建DOT和PDF文件
    create_dot_from_json(json_data, os.path.abspath(output))
    
    # 创建增强的JSON文件
    create_enhanced_fsm_json(json_data, os.path.abspath(output))



# python ./json2fsm.py -i ./save/test_3.4.1/get_fsm_from_llm.json -o ./save/test_3.4.1/dots
# python ./json2fsm.py -i ./save/wolfssl_3.4.1/get_fsm_from_llm.json -o ./save/wolfssl_3.4.1/dots
# python ./json2fsm2.py -i ./save/wolfssl_3.4.1/get_fsm_from_llm.json -o ./save/wolfssl_3.4.1/dots -c .\config\wolfssl.json
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='json to fsm')
    parser.add_argument('-i', '--input', help='Input config JSON file path', required=True)
    parser.add_argument('-o', '--output', help='Output dir', required=True)
    parser.add_argument('-c', '--config', help='Config file path', required=True)


    args = parser.parse_args()

    with open(args.config, encoding='utf-8') as f:
        config = json.load(f)


    main(args.input, args.output, config)
